from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
from modelscope import snapshot_download
from transformers import AutoImageProcessor
import torch

# 模型存储路径
model_dir='/root/autodl-tmp/Qwen/Qwen2.5-VL-7B-Instruct'

# bfloat16比默认float32占用更少的内存,同时保留了足够的精度以维持模型性能
# FlashAttention2，这是一种优化技术，旨在加速Transformer模型中的自注意力层计算，并降低其显存需求
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    model_dir,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    device_map="auto",
)

# 初始化处理器
# 使用最小和最大的像素范围，避免输入图像过大导致内存占用过大
min_pixels = 320 * 240
max_pixels = 320 * 240
processor = AutoProcessor.from_pretrained(model_dir, min_pixels=min_pixels, max_pixels=max_pixels)
#processor = AutoProcessor.from_pretrained(model_dir)

# max_pixels 定义了处理过程中视频帧的最大像素数限制，这有助于控制输入数据的大小
# fps 每秒只抽取一帧图像作为输入的一部分
messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "video",
                "video": "/root/autodl-tmp/video/alalei.mp4",
                "max_pixels": 320 * 240,
                "fps": 0.5,
            },
            {"type": "text", "text": "请你描述一下这部动漫剧里面的主要人物和故事"},
        ],
    }
]

# 把文本输入转成适合机器的文本
text = processor.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)

# 处理视觉信息
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
    text=[text],
    images=image_inputs,
    videos=video_inputs,
    padding=True,
    return_tensors="pt",
)

# 把输入放到GPU里面去
inputs = inputs.to("cuda")

# 模型推理
# max_new_tokens 限制了输出的文本长度
generated_ids = model.generate(**inputs, max_new_tokens=256)
generated_ids_trimmed = [
    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
output_text = processor.batch_decode(
    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print(output_text)
# 释放内存
torch.cuda.empty_cache()